package bigdata.backend.meta.utils;

import bigdata.backend.meta.dto.ColumnLineage;
import bigdata.backend.meta.dto.TableLineage;
import gudusoft.gsqlparser.EDbVendor;
import gudusoft.gsqlparser.dlineage.DataFlowAnalyzer;
import gudusoft.gsqlparser.dlineage.dataflow.model.Option;
import gudusoft.gsqlparser.dlineage.dataflow.model.xml.*;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.util.TablesNamesFinder;

import java.util.*;

public class TableAndColumnLineage {

    /**
     * 获取目标表字段 -> 原始来源字段
     */
    public static List<ColumnLineage> getColumnLineage(String targetTable, String sql) throws Exception {

        // ---- 配置 ----
        Option option = new Option();
        option.setVendor(EDbVendor.dbvhive);
        option.filterRelationTypes("fdd,fddi,frd,fdr".split(","));
        option.setLinkOrphanColumnToFirstTable(true);

        DataFlowAnalyzer analyzer = new DataFlowAnalyzer(sql, option);
        analyzer.generateDataFlow(true);

        dataflow df = analyzer.getDataFlow();

        HashMap<String, String> columnMap = new HashMap<>();
        HashMap<String, List<String>> relMap = new HashMap<>();
        for (relationship rel : df.getRelationships()) {
            if (rel == null || rel.getTarget() == null) continue;

            targetColumn tgt = rel.getTarget();

            if ("RelationRows".equalsIgnoreCase(tgt.getColumn())) continue;

            String tgtKey = key(tgt.getParent_name(), tgt.getColumn());
            if (!columnMap.containsKey(tgt.getId())) {
                columnMap.put(tgt.getId(), tgtKey);
            }

            for (sourceColumn src : rel.getSources()) {
                if ("RelationRows".equalsIgnoreCase(src.getColumn())) continue;
                String srcKey = key(src.getParent_name(), src.getColumn());
                if (!columnMap.containsKey(src.getId())) {
                    columnMap.put(src.getId(), srcKey);
                }
                if (relMap.containsKey(tgt.getId())) {
                    List<String> list = relMap.get(tgt.getId());
                    list.add(src.getId());
                } else {
                    relMap.put(tgt.getId(), new ArrayList<>(Arrays.asList(src.getId())));
                }
            }
        }

        // System.out.println("relMap: " + relMap);
        // System.out.println("columnMap: " + columnMap);
        ArrayList<ColumnLineage> columnLineages = new ArrayList<>();
        for (Map.Entry<String, String> entry : columnMap.entrySet()) {
            String id = entry.getKey();
            String fieldName = entry.getValue();

            if (fieldName.startsWith(targetTable + ".")) {
                // System.out.println(id + ":" + fieldName);

                // ----------- 递归展开最终来源字段 --------------
                List<String> finalSources = new ArrayList<>();
                expandLineage(id, relMap, columnMap, finalSources);
                // System.out.println("最终血缘 " + fieldName + " <- " + finalSources);
                ColumnLineage columnLineage = new ColumnLineage(fieldName, finalSources);
                columnLineages.add(columnLineage);
            }
        }
        return columnLineages;
    }

    /**
     * 递归展开血缘（id -> 最终来源字段列表）
     */
    private static void expandLineage(String id, Map<String, List<String>> relMap, Map<String, String> columnMap, List<String> result) {

        // 如果此 id 无子来源 = 基础字段
        if (!relMap.containsKey(id)) {
            result.add(columnMap.get(id));
            return;
        }

        // 有子来源，继续展开
        for (String sid : relMap.get(id)) {
            expandLineage(sid, relMap, columnMap, result);
        }
    }

    private static String key(String parent, String col) {
        if (parent == null || parent.isEmpty()) return col;
        return parent + "." + col;
    }

    /**
     * 解析SQL语句，提取目标表和源表信息
     * @param sql 待解析的SQL语句
     * @return TableLineage POJO对象，包含目标表和源表列表
     * @throws JSQLParserException 如果SQL解析失败
     */
    public static TableLineage getTableLineage(String sql) throws JSQLParserException {
        // 预处理SQL
        // 替换OVERWRITE TABLE
        String processedSql = sql.replaceAll("(?i)OVERWRITE\\s+TABLE", "INTO");
        // 替换 PARTITION语法
        processedSql = processedSql.replaceAll("(?i)\\s+PARTITION\\s*\\([^)]*\\)", "");

        // 使用TablesNamesFinder获取所有表名
        Set<String> allTableNames = TablesNamesFinder.findTables(processedSql);
        // System.out.println("All tables found: " + allTableNames);

        // 解析SQL语句
        Statement statement = CCJSqlParserUtil.parse(processedSql);

        if (!(statement instanceof Insert)) {
            throw new IllegalArgumentException("Only INSERT statements are supported");
        }

        Insert insert = (Insert) statement;
        String targetTable = insert.getTable().getFullyQualifiedName();
        Set<String> allTables = TablesNamesFinder.findTables(processedSql);
        allTables.remove(targetTable); // 移除目标表，剩余为源表

        // 创建并返回POJO对象
        return new TableLineage(targetTable, new ArrayList<>(allTables));
    }
}
